package com.gitee.mp.core.connector;

import com.gitee.mp.configure.MilvusConfigure;
import com.gitee.mp.extend.MilvusConfigureChain;
import com.gitee.mp.configure.MilvusConnection;
import com.gitee.mp.infra.constants.MilvusPlusConstants;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author hongda.li
 */
@Slf4j
public class MilvusConnectorFactory {

    @Getter
    private final MilvusConfigure configure;

    private final Map<String, MilvusConnector> connectorMap = new ConcurrentHashMap<>();

    public MilvusConnectorFactory(List<MilvusConfigureChain> configureChains) {
        this.configure = new MilvusConfigure();
        configureChains.forEach(chain -> chain.doChain(configure));
        this.configure.getConnectionList().forEach(this::registerConnector);
    }

    /**
     * 获取默认的 Milvus 连接器
     *
     * @return Milvus 连接器
     */
    public MilvusConnector getConnector() {
        return this.getConnector(MilvusPlusConstants.DEFAULT_ID);
    }

    /**
     * 获取 Milvus 连接器
     *
     * @param id 连接器 ID
     * @return Milvus 连接器
     */
    public MilvusConnector getConnector(String id) {
        return this.connectorMap.get(id);
    }

    /**
     * 移除 Milvus 连接器
     *
     * @param id 连接器 ID
     * @return Milvus 连接器
     */
    @SuppressWarnings("UnusedReturnValue")
    public MilvusConnector removeConnector(String id) {
        MilvusConnector remove = this.connectorMap.remove(id);
        Boolean success = Optional.ofNullable(remove).map(MilvusConnector::remove).orElse(true);
        if (Boolean.TRUE.equals(success)) {
            log.debug("Remove connector success : {}", id);
        } else {
            log.error("Remove connector failure : {}", id);
        }
        return remove;
    }

    /**
     * 注册 Milvus 连接器
     *
     * @param connection 配置信息
     * @return Milvus 连接器
     */
    public MilvusConnector registerConnector(MilvusConnection connection) {
        String connectorId = connection.getConnectorId();

        // 移除可能存在的连接器
        this.removeConnector(connectorId);

        MilvusConnector connector = new MilvusConnector(connection);
        this.connectorMap.put(connectorId, connector);

        log.debug("Register connector success: {}", connectorId);

        return connector;
    }
}
